import asyncio
import json
import os

from typing import Optional
from contextlib import AsyncExitStack

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()


llm_client = OpenAI(
    base_url=os.getenv("API_URL"),
    api_key=os.getenv("OPENAI_API_KEY"),
)


class MCPClient:
    def __init__(self):
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()

    async def process_query(self, query: str) -> str:
        messages = [
            {
                "role": "user",
                "content": query
            }
        ]

        response = await self.session.list_tools()
        available_tools = [{
            "type": "function",
            "function": {
                "name": tool.name,
                "description": tool.description,
                "parameters": {
                    "type": tool.inputSchema["type"],
                    "required": tool.inputSchema["required"],
                    "properties": tool.inputSchema["properties"],
                }
            }
        } for tool in response.tools]
        print(json.dumps(available_tools, indent=4))

        # Initial Claude API call
        first_response = llm_client.chat.completions.create(
            model=os.getenv("MODEL_NAME"),
            messages=messages,
            tools=available_tools,
            tool_choice="auto",
            max_tokens=4096,
            temperature=0,
        )

        messages.append(
            {
                "role": "assistant",
                "content": first_response.choices[0].message.content,
                "tool_calls": first_response.choices[0].message.tool_calls,
            }
        )

        stop_reason = (
            "tool_calls"
            if first_response.choices[0].message.tool_calls is not None
            else first_response.choices[0].finish_reason
        )

        if stop_reason == "tool_calls":
            # Extract tool use details from response
            for tool_call in first_response.choices[0].message.tool_calls:
                arguments = (
                    json.loads(tool_call.function.arguments)
                    if isinstance(tool_call.function.arguments, str)
                    else tool_call.function.arguments
                )
                # Call the tool with the arguments using our callable initialized in the tools dict
                tool_result = await self.session.call_tool(tool_call.function.name, arguments=arguments)
                # Add the tool result to the messages list
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "name": tool_call.function.name,
                        "content": tool_result.content[0].text,
                    }
                )

            # Query LLM with the user query and the tool results
            new_response = llm_client.chat.completions.create(
                model=os.getenv("MODEL_NAME"),
                messages=messages,
            )

        elif stop_reason == "stop":
            # If the LLM stopped on its own, use the first response
            new_response = first_response

        else:
            raise ValueError(f"Unknown stop reason: {stop_reason}")

            # Add the LLM response to the messages list
        messages.append(
            {"role": "assistant", "content": new_response.choices[0].message.content}
        )

        # Return the LLM response and messages
        return new_response.choices[0].message.content


    async def chat_loop(self):
        """Run an interactive chat loop"""
        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")

        while True:
            try:
                query = input("\nQuery: ").strip()

                if query.lower() == 'quit':
                    break

                response = await self.process_query(query)
                print("\n" + response)

            except Exception as e:
                print(f"\nError: {str(e)}")

    async def cleanup(self):
        """Clean up resources"""
        await self.exit_stack.aclose()


    async def connect_to_server(self, server_script_path: str):
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')
        if not (is_python or is_js):
            raise ValueError("Server script must be a .py or .js file")

        command = "python" if is_python else "node"
        server_params = StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None
        )

        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))

        await self.session.initialize()

        tools_response = await self.session.list_tools()
        tools = tools_response.tools
        # print("Connected to server with tools:", [tool for tool in tools])

        prompts_response = await self.session.list_prompts()
        prompts = prompts_response.prompts
        # print("Connected to server with prompts:", [prompt for prompt in prompts])

        resources_templates_response = await self.session.list_resource_templates()
        resources = resources_templates_response.resourceTemplates
        # print("Connected to server with resources:", [resource for resource in resources])


async def main():
    client = MCPClient()
    try:
        await client.connect_to_server("./server.py")
        await client.chat_loop()
    finally:
        await client.cleanup()


if __name__ == "__main__":
    asyncio.run(main())